package com.shanks.framework.mq.queue;

import com.shanks.framework.mq.enums.MessageRecordStatusEnum;
import com.shanks.framework.mq.enums.QueueStatusEnum;
import com.shanks.framework.mq.model.MessageRecordPO;
import jakarta.annotation.PostConstruct;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 队列管理
 */
@Component
public class MessageQueueManager {
    private static final Logger logger = LoggerFactory.getLogger(MessageQueueManager.class);

    @Value("${message-transmit-pending.queue.size:10000}")
    private int queueSize;

    private final Map<MessageRecordStatusEnum, BlockingQueue<MessageRecordPO>> messageStatus2Queue = new HashMap<>();


    @PostConstruct
    public void init() {
        // Initialize the queues for different message statuses
        /**
         * 待发送队列
         */
        messageStatus2Queue.put(MessageRecordStatusEnum.PENDING, new ArrayBlockingQueue<>(queueSize));

        /**
         * 发送中队列
         */
        messageStatus2Queue.put(MessageRecordStatusEnum.SENDING, new ArrayBlockingQueue<>(queueSize));

        /**
         * 已发送队列
         */
        messageStatus2Queue.put(MessageRecordStatusEnum.SENT, new ArrayBlockingQueue<>(queueSize));
    }

    /**
     * Batch offer messages to the respective queue.
     */
    public void offerBatchMessages(List<MessageRecordPO> messageRecords, MessageRecordStatusEnum messageRecordStatus) {
        for (MessageRecordPO record : messageRecords) {
            // 入队失败
            if (QueueStatusEnum.FAIL == offerMessageIfAbsent(record, messageRecordStatus)) {
                logger.warn("Failed to add batch of messages to the queue, Queue is full. Current message: {}, Batch message IDs: {}",
                        record, messageRecords.stream()
                                .map(MessageRecordPO::getMessageId)
                                .collect(Collectors.toList()));
                break;
            }
        }
    }

    /**
     * Offer a message to the respective queue if it's not already present.
     */
    public QueueStatusEnum offerMessageIfAbsent(MessageRecordPO messageRecord, MessageRecordStatusEnum status) {

        if (isMessageAlreadyInQueue(messageRecord)) {
            return QueueStatusEnum.EXIST;
        }

        boolean resultFlag = false;
        try {
            BlockingQueue<MessageRecordPO> queue = messageStatus2Queue.get(status);
            if (queue != null) {
                resultFlag = queue.offer(messageRecord, 1L, TimeUnit.SECONDS);
            }
        } catch (InterruptedException e) {
            logger.error("Failed to add message to queue, Message: {}", messageRecord, e);
        }

        return resultFlag ? QueueStatusEnum.SUCCESS : QueueStatusEnum.FAIL;
    }

    /**
     * Take a message from the respective queue based on the status.
     */
    public MessageRecordPO takeMessage(MessageRecordStatusEnum status) throws InterruptedException {
        BlockingQueue<MessageRecordPO> queue = messageStatus2Queue.get(status);
        return queue != null ? queue.take() : null;
    }

    /**
     * Check if the message ID is already present in any queue.
     */
    private boolean isMessageAlreadyInQueue(MessageRecordPO messageRecordPO) {
        return messageStatus2Queue.values().contains(messageRecordPO);
    }

    /**
     * Get the size of the respective queue.
     */
    public int getQueueSize(MessageRecordStatusEnum status) {
        BlockingQueue<MessageRecordPO> queue = messageStatus2Queue.get(status);
        return queue != null ? queue.size() : -1;
    }


}
